import requests
from kimi.KimiToken import new_token
from logger import logger


class KimiChatSession:
    def __init__(self):
        self.access_token = "" # 生成唯一的token
        self.messages = []  # 用于存储会话中的消息
        # self.authorization = 'Bearer ' + self.access_token

        self.base_url = "https://kimi.moonshot.cn/api/chat"
        self.chat_id = self.new_chat_id()

    def get_base_headers(self):
        # if self.access_token is None:
        #     self.access_token = new_token()[0]
        headers = {
            "accept": "*/*",
            "accept-language": "zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6",
            "authorization": 'Bearer ' + self.access_token,
            "content-type": "application/json",
            "r-timezone": "Asia/Shanghai",
            "sec-ch-ua": "\"Not A(Brand\";v=\"99\", \"Microsoft Edge\";v=\"121\", \"Chromium\";v=\"121\"",
            "sec-ch-ua-mobile": "?0",
            "sec-ch-ua-platform": "\"Windows\"",
            "sec-fetch-dest": "empty",
            "sec-fetch-mode": "cors",
            "sec-fetch-site": "same-origin"
        }
        return headers;


    def start_session(self):
        logger.info(f"欢迎进入kimi会话:")
        self.interactive_chat()

    def interactive_chat(self):
        while True:
            user_input = input("请输入消息（输入'1'结束会话）：")
            if user_input.lower() == '1':
                self.end_session()
                break
            self.messages.append(user_input)
            print(self.chat(user_input))

    def get_messages(self):
        return self.messages

    def end_session(self):
        self.del_chat()
        print(f"会话历史：")
        for message in self.messages:
            print(message)

    def new_chat_id(self):
        data = {
            "name": "KimiChat Stock info",
            "is_example": False
        }
        r = requests.post(self.base_url, headers=self.get_base_headers(), json=data)
        if r.status_code == 200:
            chatid = r.json()['id']
            logger.info(chatid)
            return chatid
        elif r.status_code == 401:
            self.access_token = new_token()[0]
            return self.new_chat_id()
        else:
            exit('chat_id get error!')
            return None

    def chat(self,text):
        data2 = {
            "messages": [
                {
                    "role": "user",
                    "content": text
                }
            ],
            "refs": [],
            "use_search": True
        }
        url2 = 'https://kimi.moonshot.cn/api/chat/' + str(self.chat_id) + '/completion/stream'
        r2 = requests.post(url=url2, json=data2, headers=self.get_base_headers())
        if r2.status_code == 200:
            url3 = 'https://kimi.moonshot.cn/api/chat/' + str(self.chat_id) + '/segment/scroll'
            headers3 = {
                "accept": "*/*",
                "accept-language": "zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6",
                "authorization": 'Bearer ' + self.access_token,
                "content-type": "application/json",
                "r-timezone": "Asia/Shanghai",
                "sec-ch-ua": "\"Not A(Brand\";v=\"99\", \"Microsoft Edge\";v=\"121\", \"Chromium\";v=\"121\"",
                "sec-ch-ua-mobile": "?0",
                "sec-ch-ua-platform": "\"Windows\"",
                "sec-fetch-dest": "empty",
                "sec-fetch-mode": "cors",
                "sec-fetch-site": "same-origin",
                ## "referrer": "https://kimi.moonshot.cn/chat/cn85ml83qff4rjdhi2r0?data_source=tracer&track_id=pbaes.Pi_gb2BTHE2MNxEl7YxYIbDQyqJz5l-ly9iBqnFyBZ8aCfIjAMbPj-PIFU6A7lS3Fw4NLkI3U52eof9gDi6kIU_MjTd5FNDPNyvcdQECQIUHVHAyKTVfJe37Nzy16BXqyYTtR0RqE4SM9lF2OJOR-IxDS4fhD82jsGYCaHc81hCRfURp_sm1E81E5IztMgbd&utm_campaign=TR_ikmtSa64&utm_content=&utm_medium=B%E7%AB%A0%E7%AB%99PC%E7%AB%AF%E5%B9%BF%E5%91%8A&utm_source=bilibili&utm_term=",
                "referrer": "https://kimi.moonshot.cn/chat/" + str(self.chat_id),
                "referrerPolicy": "strict-origin-when-cross-origin"
            }
            data3 = {
                "last": 50
            }

            r3 = requests.post(url3, headers=headers3, json=data3)
            if r3.status_code == 200:
                result = r3.json()
                # print(result)
                return self.get_content(result)
            else:
                exit('ERROR_Number:3')
                return None
        elif r2.status_code == 401:
            self.access_token = new_token()[0]
            return self.chat(text)
        else:
            exit('ERROR_Number:2')
            return None

    def del_chat(self):
        url4 = "https://kimi.moonshot.cn/api/chat/" + str(self.chat_id)
        headers4 = {
            "accept": "*/*",
            "accept-language": "zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6",
            "authorization": 'Bearer ' + self.access_token,
            "content-type": "application/json",
            "r-timezone": "Asia/Shanghai",
            "sec-ch-ua": "\"Not A(Brand\";v=\"99\", \"Microsoft Edge\";v=\"121\", \"Chromium\";v=\"121\"",
            "sec-ch-ua-mobile": "?0",
            "sec-ch-ua-platform": "\"Windows\"",
            "sec-fetch-dest": "empty",
            "sec-fetch-mode": "cors",
            "sec-fetch-site": "same-origin"
        }

        r4 = requests.delete(url4, headers=headers4)
        if r4.status_code == 200:
            print("会话："+self.chat_id+"删除成功！")
        else:
            exit('ERROR_Number:4')
            return None

    def get_content(self,data):
        content = data['items'][-1]['content'] +"\n"
        _fist = True
        # 遍历每个元素中的 search_plus 列表
        for item in data['items']:
            if 'search_plus' in item:
                for msg in item['search_plus']:
                    if 'msg' in msg and msg['msg']['type'] == 'get_res':
                        if _fist:
                            content += "-------------url------------\n"
                            _fist = False
                        title = msg['msg']['title']
                        url = msg['msg']['url']
                        content += f"Title: {title}, URL: {url}\n"
        return content


class KimiStockSession(KimiChatSession):
    def __init__(self):
        super().__init__()  # 调用基类的构造函数
        self.stock_init = "给你一个股票的走势图数据，包含date,code,open,high,low,close,preclose,volume,amount，分别表示交易日，股票代码，开盘价，最高价，最低价，收盘价，前一日收盘价，成交量，成交金额，请你用日本蜡烛图、趋势理论、短信交易理论分别预测后期的走势，以及下一个交易日购买交易策略，并说明理由。你理解我的需求么？理解输出'ok'，否则输出'No'。"  # 股票特有的token


    def init_stock_purchase_decision(self):
        rs = self.chat(self.stock_init)
        print(rs)
        return rs.lower().startswith("ok")

    def stock_message(self,code,data):
        stock_data_string = f"股票代码：{code} 数据如下：\n"
        for row in data:
            stock_data_string += ",".join(map(str, row)) + "\n"
        return stock_data_string

if __name__ == '__main__':
    data = {
    "items": [
        {
            "id": "cnab7vmcp7fevuraoqug",
            "content": "学习python有前途么？",
            "context_type": "text",
            "created_at": "2024-02-20T22:15:26.952336+08:00",
            "role": "user",
            "group_id": "cnab7vmcp7fevuraoqu0"
        },
        {
            "id": "cnab7vmcp7fevuraoqv0",
            "content": " 学习Pytho发展机会。",
            "context_type": "text",
            "created_at": "2024-02-20T22:15:26.952336+08:00",
            "role": "assistant",
            "search_plus": [
                {
                    "event": "search_plus",
                    "msg": {
                        "type": "start"
                    }
                },
                {
                    "event": "search_plus",
                    "msg": {
                        "type": "start_res"
                    }
                },
                {
                    "event": "search_plus",
                    "msg": {
                        "successNum": 1,
                        "title": "Python 就业前景如何? - 知乎",
                        "type": "get_res",
                        "url": "https://www.zhihu.com/question/618574787"
                    }
                },
                {
                    "event": "search_plus",
                    "msg": {
                        "successNum": 2,
                        "title": "学习Python之后的七大就业方向 - 知乎",
                        "type": "get_res",
                        "url": "https://zhuanlan.zhihu.com/p/495287433"
                    }
                },
                {
                    "event": "search_plus",
                    "msg": {
                        "successNum": 3,
                        "title": "未来十年Python的前景会怎样？ - 知乎",
                        "type": "get_res",
                        "url": "https://www.zhihu.com/question/22112542"
                    }
                },
                {
                    "event": "search_plus",
                    "msg": {
                        "successNum": 4,
                        "title": "清华教授：分析Python未来的前景(文末附福利) - 知乎",
                        "type": "get_res",
                        "url": "https://zhuanlan.zhihu.com/p/152156350"
                    }
                },
                {
                    "event": "search_plus",
                    "msg": {
                        "successNum": 5,
                        "title": "Python发展有前途吗？ - 知乎",
                        "type": "get_res",
                        "url": "https://www.zhihu.com/question/434361361"
                    }
                },
                {
                    "event": "search_plus",
                    "msg": {
                        "successNum": 6,
                        "title": "最全的python就业前景分析以及自学建议 - 知乎",
                        "type": "get_res",
                        "url": "https://zhuanlan.zhihu.com/p/343915998"
                    }
                },
                {
                    "event": "search_plus",
                    "msg": {
                        "type": "done",
                        "url_list": [
                            {
                                "title": "Python发展有前途吗？ - 知乎",
                                "url": "https://www.zhihu.com/question/434361361"
                            },
                            {
                                "title": "最全的python就业前景分析以及自学建议 - 知乎",
                                "url": "https://zhuanlan.zhihu.com/p/343915998"
                            },
                            {
                                "title": "Python 就业前景如何? - 知乎",
                                "url": "https://www.zhihu.com/question/618574787"
                            },
                            {
                                "title": "清华教授：分析Python未来的前景(文末附福利) - 知乎",
                                "url": "https://zhuanlan.zhihu.com/p/152156350"
                            },
                            {
                                "title": "学习Python之后的七大就业方向 - 知乎",
                                "url": "https://zhuanlan.zhihu.com/p/495287433"
                            },
                            {
                                "title": "未来十年Python的前景会怎样？ - 知乎",
                                "url": "https://www.zhihu.com/question/22112542"
                            }
                        ]
                    }
                },
                {
                    "event": "search_plus",
                    "msg": {
                        "type": "answer"
                    }
                }
            ],
            "group_id": "cnab7vmcp7fevuraoqu0"
        }
    ],
    "total": 2
}

    kimi = KimiChatSession()
    print(kimi.get_content(data))
